import os  
from dotenv import load_dotenv  
load_dotenv()  
from llama_index.core.agent import ReActAgent  
from llama_index.llms.ollama import Ollama  
from llama_index.core.tools import FunctionTool  
import pexpect  
import time  
import networkx as nx
import re
import numpy as np

# 定义案例字典，键为案例名称，值为相关描述  
case_files_dict = {  
        "case9.m": "Smart grid data sample case9.m"
        
}  

# 创建一个模板  
question_host = "Analyze if there's an overload in the grid sample: {case_file},answer how many overloaded branches there are give me their overloaded branches and judge whether there is failure cascade. "  
def matpower_check_overload(case_name: str) -> str:
    """通过 MATPOWER 执行潮流计算并检查过载支路"""
    import pexpect
    import re
    import numpy as np

    # 启动 Docker 并进入 Octave 环境
    child = pexpect.spawn('docker run -it --rm matpower/matpower:latest octave-cli')
    child.expect('octave:', timeout=300)

    try:
        # ========== 加载案例和计算潮流 ========== 
        child.sendline(f"mpc = loadcase('{case_name}')")
        child.sendline("results = runpf(mpc)")
        child.expect('octave:', timeout=300)

        # ========== 获取关键数据 ========== 
        child.sendline("""
            branch = results.branch;
            S_from = sqrt(branch(:,14).^2 + branch(:,15).^2);
            S_to = sqrt(branch(:,16).^2 + branch(:,17).^2);
            fprintf('====BEGIN_DEBUG====\\n');
            disp(branch(:,6));   % Rate_A
            disp(S_from);
            disp(S_to);
            fprintf('====END_DEBUG====\\n');
        """)
        child.expect('====BEGIN_DEBUG====', timeout=30)
        child.expect('====END_DEBUG====', timeout=30)
        debug_data = child.before.decode('utf-8')

        # ========== 核心修改：科学计数法转换 ========== 
        def convert_sci(num_str):
            """将科学计数法字符串转为浮点数"""
            if 'e' in num_str.lower():
                return float(num_str)
            return float(num_str)  # 普通数字直接转换

        # 提取所有数值（支持多列数据和科学计数法） 
        nums = []
        for line in debug_data.splitlines():
            line = line.strip()
            if line:
                # 分割每行的多个数值
                parts = re.split(r'\s+', line)
                for part in parts:
                    if re.match(r'^[+-]?\d*\.?\d+([eE][+-]?\d+)?$', part):
                        nums.append(convert_sci(part))

        # 分组校验 
        if len(nums)%3 != 0:
            raise ValueError("数据解析错误，请检查输出格式")
        
        group_len = len(nums)//3
        rate_a = np.array(nums[:group_len])
        s_from = np.array(nums[group_len:2*group_len])
        s_to = np.array(nums[2*group_len:3*group_len])

        # ========== 打印详细的 S_from, S_to 和 Rate_A ==========
        print("编号   Rate_A    S_from    S_to")
        for i, (r, s_f, s_t) in enumerate(zip(rate_a, s_from, s_to), 1):
            print(f"{i:<5} {r:.6f}    {s_f:.6f}    {s_t:.6f}")

        # ========== 简化过载判断 ========== 
        overloads = []
        for i in range(group_len):
            if rate_a[i] <= 0:
                continue  # 忽略无容量限制的支路
            
            # 直接数值比较（考虑浮点精度）
            if (s_from[i] > rate_a[i] + 1e-3) or (s_to[i] > rate_a[i] + 1e-3):
                overloads.append(i+1)  # MATLAB索引从1开始

        # ========== 结果生成 ========== 
        result = (
            f"过载支路分析结果\n"
            f"过载支路数量: {len(overloads)}\n"
            f"过载支路编号: {overloads}"
        )

    finally:
        child.sendline('exit')
        child.close()

    return result

def check_failure_cascade(case_name: str) -> bool:
    """检查电网案例是否发生失败级联现象"""
    print(f"[DEBUG] 启动 Docker 容器并运行 Octave...")
    child = pexpect.spawn('docker run -it --rm -e LANG=en_US.UTF-8 -e TERM=dumb matpower/matpower:latest octave-cli')
    child.expect(r'octave.*> ', timeout=900)
    
    print(f"[DEBUG] 加载案例 {case_name} ...")
    child.sendline(f"mpc = loadcase('{case_name}');")
    child.expect(r'octave.*> ', timeout=900)

    print(f"[DEBUG] 运行潮流计算...")
    child.sendline("results = runpf(mpc);")
    child.expect(r'octave.*> ', timeout=600)

    # ========== 获取关键数据 ========== 
    child.sendline("""
        branch = results.branch;
        S_from = sqrt(branch(:, 14).^2 + branch(:, 15).^2);
        S_to = sqrt(branch(:, 16).^2 + branch(:, 17).^2);
        RATE_A = branch(:, 6);
        fprintf('====BEGIN_DEBUG====\\n');
        disp(RATE_A);   % Rate_A
        disp(S_from);   % S_from
        disp(S_to);     % S_to
        fprintf('====END_DEBUG====\\n');
    """)
    child.expect('====BEGIN_DEBUG====', timeout=30)
    child.expect('====END_DEBUG====', timeout=30)
    debug_data = child.before.decode('utf-8')

    def convert_sci(num_str):
        if 'e' in num_str.lower():
            return float(num_str)
        return float(num_str)

    nums = []
    for line in debug_data.splitlines():
        line = line.strip()
        if line:
            parts = re.split(r'\s+', line)
            for part in parts:
                if re.match(r'^[+-]?\d*\.?\d+([eE][+-]?\d+)?$', part):
                    nums.append(convert_sci(part))

    if len(nums) % 3 != 0:
        raise ValueError("数据解析错误，请检查输出格式")
    
    group_len = len(nums) // 3
    rate_a = np.array(nums[:group_len])
    s_from = np.array(nums[group_len:2 * group_len])
    s_to = np.array(nums[2 * group_len:3 * group_len])

    print("编号   Rate_A    S_from    S_to")
    for i, (r, s_f, s_t) in enumerate(zip(rate_a, s_from, s_to), 1):
        print(f"{i:<5} {r:.6f}    {s_f:.6f}    {s_t:.6f}")

    # ========== 严格判断过载条件 ==========
    overloads = []
    for i in range(group_len):
        if rate_a[i] <= 0:
            continue
        tolerance = 1e-3 + 1e-3 * rate_a[i]
        if (s_from[i] > rate_a[i] + tolerance) or (s_to[i] > rate_a[i] + tolerance):
            overloads.append(i + 1)

    print(f"[DEBUG] 初始过载支路: {overloads}")

    # ========== 迭代处理过载支路 ==========
    iteration = 0
    max_iterations = 10
    removed_branches = set()

    while overloads and iteration < max_iterations:
        iteration += 1
        print(f"[DEBUG] 迭代次数: {iteration}")
        current_overloads = [idx for idx in overloads if idx not in removed_branches]
        if not current_overloads:
            break

        branches_str = '[' + ','.join(map(str, current_overloads)) + ']'
        child.sendline(f"mpc.branch({branches_str}, :) = [];")
        child.expect(r'octave.*> ', timeout=300)
        removed_branches.update(current_overloads)

        # 检查潮流计算是否成功
        child.sendline("try; results = runpf(mpc); catch; fprintf('FAILED\\n'); end;")
        child.expect(r'octave.*> ', timeout=600)
        if 'FAILED' in child.before.decode('utf-8'):
            print("[DEBUG] 潮流计算不收敛，终止迭代")
            break

        # 获取新的过载支路
        child.sendline("""
            try
                branch = results.branch;
                S_from = sqrt(branch(:, 14).^2 + branch(:, 15).^2);
                S_to = sqrt(branch(:, 16).^2 + branch(:, 17).^2);
                RATE_A = branch(:, 6);
                overloaded_branches = find(S_from > RATE_A | S_to > RATE_A);
                fprintf('====BEGIN_OVERLOADS====\\n');
                disp(overloaded_branches);
                fprintf('====END_OVERLOADS====\\n');
            catch
                fprintf('====BEGIN_OVERLOADS====\\n[]\\n====END_OVERLOADS====\\n');
            end
        """)
        child.expect('====BEGIN_OVERLOADS====', timeout=30)
        child.expect('====END_OVERLOADS====', timeout=30)
        overload_output = child.before.decode('utf-8')

        # 安全提取数字
        overloads = []
        matches = re.findall(r'\b\d+\b', overload_output)
        if matches:
            overloads = list(map(int, matches))
        print(f"[DEBUG] 新的过载支路: {overloads}")

    # ========== 孤岛分析 ========== 
    child.sendline("try; disp(size(results.branch, 1)); catch; disp(0); end;")
    child.expect(r'octave.*> ', timeout=300)
    branch_count = 0
    match = re.search(r'\d+', child.before.decode('utf-8'))
    if match:
        branch_count = int(match.group())

    num_islands = 1
    if branch_count > 0:
        child.sendline("""
            try
                [groups, isolated] = find_islands(mpc);
                fprintf('====BEGIN_ISLANDS====\\n%d\\n====END_ISLANDS====\\n', length(groups));
            catch
                fprintf('====BEGIN_ISLANDS====\\n1\\n====END_ISLANDS====\\n');
            end
        """)
        child.expect('====BEGIN_ISLANDS====', timeout=30)
        child.expect('====END_ISLANDS====', timeout=30)
        islands_output = child.before.decode('utf-8')
        match = re.search(r'\d+', islands_output)
        if match:
            num_islands = int(match.group())

    print(f"[DEBUG] 孤岛数量: {num_islands}")
    
    child.sendline('exit')
    child.close()
    return num_islands > 1


def main():  
    action_tool_1 = FunctionTool.from_defaults(fn=matpower_check_overload)  
    action_tool_2 = FunctionTool.from_defaults(fn=check_failure_cascade)  
    llm = Ollama(model="Gemma3:4b", request_timeout=720.0)  

    output = ""  

    # 遍历案例字典  
    for case_name, policy in case_files_dict.items():  
        print(f"Processing {case_name}...")  

        question = question_host.format(case_file=case_name)  
        context_question = f"Given the context that: \"{policy}\"\n{question}"  
        agent = ReActAgent.from_tools([action_tool_1,action_tool_2], llm=llm, verbose=True,max_iterations=100)  
        response = agent.chat(context_question)  
        print("response:\n")  
        print(response)  
        output += f"{question}\n{response}\n"  

    print("output:\n")  
    print(output)  

    output_file_path = "output_results1.txt"  
    with open(output_file_path, "w") as file:  
        file.write(output)  

    print(f"Processing completed. Results are written to {output_file_path}.")  

if __name__ == "__main__":  
    main()